﻿using System;
using System.Collections.Generic;
using System.Drawing;
using System.IO;
using System.Globalization;
using System.Runtime.Serialization.Formatters.Binary;
using EvoBoost.ComputerVision.ViolaJones;
using EvoBoost.Core;
using System.Diagnostics;
using System.Text;

namespace EvoBoost.Trainer.ViolaJones
{
	class Program
	{
		private static bool TryParseDouble(string s, out double result)
		{
			return Double.TryParse(s, NumberStyles.Float, CultureInfo.InvariantCulture, out result);
		}

		private static bool TryParseInt(string s, out int result)
		{
			return Int32.TryParse(s, NumberStyles.Integer, CultureInfo.InvariantCulture, out result);
		}

		static void Main(string[] args)
		{
#if !(DEBUG)
			try
#endif
			{
				// Read parameters
				int iterations, messageLevel, maxTreeRegions, minSamplesInNode, populationSize;
				double adjustProb, mutationRate, crossoverRate, shrinkage;
				if (args.Length != 11 ||
					!TryParseInt(args[2], out iterations) || !TryParseInt(args[3], out messageLevel) ||
					!TryParseInt(args[4], out maxTreeRegions) || !TryParseInt(args[5], out minSamplesInNode) ||
					!TryParseDouble(args[6], out shrinkage) || !TryParseInt(args[7], out populationSize) ||
					!TryParseDouble(args[8], out adjustProb) || !TryParseDouble(args[9], out mutationRate) ||
					!TryParseDouble(args[10], out crossoverRate))
				{
					Console.WriteLine(
						"Usage: {0} <samplePath> <resultModelName> <iterations> <messageLevel> <maxTreeRegions> <minSamplesInNode> <shrinkage> " +
						"<populationSize> <adjustProb> <mutationRate> <crossoverRate>",
						Environment.GetCommandLineArgs()[0]);
					return;
				}

				// Create all required stuff
				ViolaJonesSampleSet trainingSet = LoadViolaJonesSampleSetFromFolder(Path.Combine(args[0], "Train"));
				ViolaJonesSampleSet testSet = LoadViolaJonesSampleSetFromFolder(Path.Combine(args[0], "Test"));
				if (trainingSet.ImageSize != testSet.ImageSize)
					throw new ApplicationException("Image sizes in training and test sets are not equal.");
				Size imageSize = trainingSet.ImageSize;
				Console.WriteLine();

				// Create and setup experiment runner
				ExperimentRunner<IntegralImage> experimentRunner = new ExperimentRunner<IntegralImage>();
				experimentRunner.Split = HaarFeatureSplit.Create(imageSize, false);
				experimentRunner.Booster.MessageLevel = messageLevel;
				experimentRunner.Booster.BoostingIterations = iterations;
				experimentRunner.Booster.AdjustSplitterMutationProbability = adjustProb;
				experimentRunner.Booster.MutationRate = mutationRate;
				experimentRunner.Booster.CrossoverRate = crossoverRate;
				experimentRunner.Booster.MaxTreeRegions = maxTreeRegions;
				experimentRunner.Booster.MinSamplesInNode = minSamplesInNode;
				experimentRunner.Booster.PopulationSize = populationSize;
				experimentRunner.Booster.Shrinkage = shrinkage;
				experimentRunner.Booster.ProxyLossTrimmingRate = 0.02;
				experimentRunner.Booster.ProxyLossFunction = new ExponentialLossFunction();
				experimentRunner.Booster.LossFunction = new MisclassificationLossFunction();

				// Run train/test experiment
				TrainingTestStatistics<IntegralImage> stats = experimentRunner.RunTrainingAndTest(trainingSet.Samples, testSet.Samples);

				// Print errors on training and test sets
				Console.WriteLine("Training errors:");
				for (int i = 0; i < stats.TrainingErrors.Count; ++i)
					Console.WriteLine("{0:0.000000}", stats.TrainingErrors[i]);
				Console.WriteLine("Test errors:");
				for (int i = 0; i < stats.TestErrors.Count; ++i)
					Console.WriteLine("{0:0.000000}", stats.TestErrors[i]);

				// Save model
				SaveModel(stats.Model, args[1]);
			}
#if !(DEBUG)
			catch (Exception e)
			{
				Console.WriteLine(e.Message);
			}
#endif
		}

		private static void SaveModel<T>(Forest<T> model, string file)
		{
			BinaryFormatter formatter = new BinaryFormatter();
			using (FileStream stream = new FileStream(file, FileMode.Create))
				formatter.Serialize(stream, model);
		}

		private static ViolaJonesSampleSet LoadViolaJonesSampleSetFromFolder(string pathToFolder)
		{
			Console.Write(string.Format("Loading samples from '{0}'... ", pathToFolder));

			Size positiveImageSize, negativeImageSize;
			List<Sample<IntegralImage>> samples = new List<Sample<IntegralImage>>();
			List<Bitmap> originalImages = new List<Bitmap>();
			LoadViolaJonesSampleSetFromFolder(
				Path.Combine(pathToFolder, "positives"), true, samples, originalImages, out positiveImageSize);
			LoadViolaJonesSampleSetFromFolder(
				Path.Combine(pathToFolder, "negatives"), false, samples, originalImages, out negativeImageSize);

			Console.WriteLine("Done");

			if (positiveImageSize != negativeImageSize)
				throw new ApplicationException("Positive and negative samples should have the same size.");

			return new ViolaJonesSampleSet(samples, positiveImageSize);
		}

		private static void MakeFeatureListForWuRehgDetector(string outputFile, Size imageSize)
		{
			List<string> featureCandidates = new List<string>();
			for (int x1 = 0; x1 < imageSize.Width; ++x1)
				for (int y1 = 0; y1 < imageSize.Height; ++y1)
					for (int x2 = x1 + 1; x2 <= imageSize.Width; ++x2)
						for (int y2 = y1 + 1; y2 <= imageSize.Height; ++y2)
						{
							int width = x2 - x1, height = y2 - y1;
							if (width % 2 == 0)
							{
								// Two stripes vertical
								featureCandidates.Add(String.Format(
									"0 0 0 {0} {1} {2} -1 {3} -1 {4} -1", x1, x1 + width / 2, x1 + width, y1, y1 + height));
							}
							if (height % 2 == 0)
							{
								// Two stripes horizontal
								featureCandidates.Add(String.Format(
									"0 0 1 {0} -1 {1} -1 {2} {3} {4} -1", x1, x1 + width, y1, y1 + height / 2, y1 + height));
							}
							if (width % 3 == 0)
							{
								// Three stripes vertical
								featureCandidates.Add(String.Format(
									"0 0 2 {0} {1} {2} {3} {4} -1 {5} -1", x1, x1 + width / 3, x1 + 2 * width / 3, x1 + width, y1, y1 + height));
							}
							if (height % 3 == 0)
							{
								// Three stripes horizontal
								featureCandidates.Add(String.Format(
									"0 0 3 {0} -1 {1} -1 {2} {3} {4} {5}", x1, x1 + width, y1, y1 + width / 3, y1 + 2 * width / 3, y1 + width));
							}
							if (width % 2 == 0 && height % 2 == 0)
							{
								// Four squares
								featureCandidates.Add(String.Format(
									"0 0 4 {0} {1} {2} -1 {3} {4} {5} -1", x1, x1 + width / 2, x1 + width, y1, y1 + height / 2, y1 + height));
							}
						}

			// Draw candidates uniformly
			const int CandidateCount = 40000;
			int currentIndex = 0;
			Random random = new Random();
			while (currentIndex < Math.Min(CandidateCount, featureCandidates.Count))
			{
				int selectedIndex = random.Next(currentIndex, featureCandidates.Count);
				string temp = featureCandidates[currentIndex];
				featureCandidates[currentIndex] = featureCandidates[selectedIndex];
				featureCandidates[selectedIndex] = temp;
				currentIndex += 1;
			}

			// Write candidates
			using (StreamWriter writer = new StreamWriter(outputFile))
			{
				for (int i = 0; i < currentIndex; ++i)
					writer.WriteLine(featureCandidates[i]);
			}
		}

		private static void PrintNumberInWuRehgFormat(BinaryWriter writer, int number)
		{
			string numberContents = number.ToString();
			foreach (char c in numberContents)
				writer.Write(c);
			writer.Write('\n');
		}

		private static void ConvertSampleSetToWuRehgFormat(string pathToFolder, string outputFile)
		{
			Console.Write(string.Format("Loading samples from '{0}'... ", pathToFolder));

			Size positiveImageSize, negativeImageSize;
			List<Sample<IntegralImage>> samples = new List<Sample<IntegralImage>>();
			List<Bitmap> originalImages = new List<Bitmap>();
			LoadViolaJonesSampleSetFromFolder(
				Path.Combine(pathToFolder, "positives"), true, samples, originalImages, out positiveImageSize);
			LoadViolaJonesSampleSetFromFolder(
				Path.Combine(pathToFolder, "negatives"), false, samples, originalImages, out negativeImageSize);

			Console.WriteLine("Done");

			if (positiveImageSize != negativeImageSize)
				throw new ApplicationException("Positive and negative samples should have the same size.");

			Console.Write("Converting... ");

			using (FileStream stream = new FileStream(outputFile, FileMode.Create, FileAccess.Write))
			using (BinaryWriter writer = new BinaryWriter(stream, Encoding.ASCII))
			{
				PrintNumberInWuRehgFormat(writer, originalImages.Count);

				for (int imageIndex = 0; imageIndex < originalImages.Count; ++imageIndex)
				{
					if (samples[imageIndex].Outcome == 1)
						PrintNumberInWuRehgFormat(writer, 1);
					else if (samples[imageIndex].Outcome == -1)
						PrintNumberInWuRehgFormat(writer, 0);
					else
						throw new ApplicationException(String.Format("Strange outcome value {0} found.", samples[imageIndex].Outcome));

					PrintNumberInWuRehgFormat(writer, positiveImageSize.Height);
					PrintNumberInWuRehgFormat(writer, positiveImageSize.Width);

					for (int i = 0; i < originalImages[imageIndex].Height; ++i)
						for (int j = 0; j < originalImages[imageIndex].Width; ++j)
						{
							int brightness = (int) (originalImages[imageIndex].GetPixel(i, j).GetBrightness() * 255);
							Debug.Assert(brightness >= 0 && brightness <= 255);
							writer.Write((byte) brightness);
						}

					writer.Write('\n');
				}
			}

			Console.WriteLine("Done");
		}

		private static void LoadViolaJonesSampleSetFromFolder(
			string pathToFolder, bool positive, ICollection<Sample<IntegralImage>> result, ICollection<Bitmap> originalImages, out Size imageSize)
		{
			Size? imageSizeDetected = null;
			foreach (string file in Directory.GetFiles(pathToFolder))
			{
				string extension = Path.GetExtension(file).ToLower();
				if (extension != ".png" && extension != ".bmp" && extension != ".jpg" && extension != ".jpeg")
					continue;

				Bitmap bitmap = new Bitmap(file);

				if (imageSizeDetected == null)
					imageSizeDetected = bitmap.Size;
				else if (imageSizeDetected != bitmap.Size)
					throw new ApplicationException("All the images in the given folder should have the same size.");

				IntegralImage image = new IntegralImage(bitmap);
				result.Add(new Sample<IntegralImage>(image, positive ? 1 : -1));
				originalImages.Add(bitmap);
			}

			if (imageSizeDetected == null)
				throw new ApplicationException(string.Format("No images were found in the directory '{0}'.", pathToFolder));

			imageSize = imageSizeDetected.GetValueOrDefault();
		}

		private class ViolaJonesSampleSet
		{
			public ICollection<Sample<IntegralImage>> Samples { get; private set; }

			public Size ImageSize { get; private set; }

			public ViolaJonesSampleSet(ICollection<Sample<IntegralImage>> samples, Size imageSize)
			{
				this.Samples = samples;
				this.ImageSize = imageSize;
			}
		}
	}
}